"""
选币策略 | 邢不行
author: 邢不行
微信: xbx6660
"""
import os
import time
import multiprocessing

from Config import *
from Functions import *
from Evaluate import *
from joblib import Parallel, delayed
pd.set_option('display.max_rows', 1000)
pd.set_option('expand_frame_repr', False)  # 当列太多时不换行

def handle_data(out_csv_path, df_, datename, factor_list, filter_list):
    df = df_.copy()
    # 构建df中参与选币的因子信息
    factor_column_list = []
    for factor_name, if_reverse, parameter_list, weight in factor_list:
        factor_column_list.append(f'{factor_name}_{str(parameter_list)}')
    # 构建df中参与过滤的因子信息
    filter_column_list = []
    for factor_name, parameter_list, if_reverse, value in filter_list:
        filter_column_list.append(f'{factor_name}_{str(parameter_list)}')
    # =合并 factor_column_list 和 filter_column_list 的信息（会去重）
    all_factor_list = list(set(factor_column_list + filter_column_list))

    # ===将数据分为现货数据和合约数据
    # 合约和现货数据分离
    df_spot = df[df['symbol_type'] == 'spot']  # 取出现货数据，后续根据配置用来进行过滤和选币
    df_swap = df[df['symbol_type'] == 'swap']  # 取出合约数据，后续根据配置用来进行过滤和选币，并替换掉现货中的行情数据

    # 是否使用现货数据进行回测，如果为True，则表明在现货数据中选币；为False，则在合约数据中选币
    if if_use_spot:  # 使用现货数据，则在现货中进行过滤，并选币
        back_test_df = df_spot.copy()
        if back_test_df.empty:  # 如果数据中不存在现货数据就退出
            print('当前为现货模式，但是数据中不存在现货数据，请重新整理数据。')
            exit()
    else:  # 使用现货数据，则在现货中进行过滤，并选币
        back_test_df = df_swap.copy()
    # 删除选币因子为空的数据
    back_test_df.dropna(subset=all_factor_list, inplace=True)

    # ===计算因子
    # 调用Strategy.calc_factor计算最终用来选币的因子，具体的因子计算方法可以在Strategy.calc_factor中自定义
    back_test_df = Strategy.calc_factor(back_test_df, external_list=factor_list)

    # =计算好因子后的一些数据整理工作
    # 获取做多和做空因子的字段名，本框架支持多头和空头是不同的因子
    long_short_columns = list({Strategy.long_factor, Strategy.short_factor})
    # 后期需要保存的列 = 多空因子字段 + 过滤因子字段
    save_columns = list(set(long_short_columns + filter_column_list))
    # 只保留一些有用的字段，减小数据大小
    back_test_df = back_test_df[['candle_begin_time', 'offset', 'symbol', 'symbol_type', 'tag', 'ret_next', '每小时涨跌幅',
                                 'fundingRate'] + save_columns]  # 只保留需要的字段
    # ===对回测数据进行过滤
    # 过滤的具体方法可以在Strategy.before_filter中自定义，故多空可能是分开过滤的，所以需要多头给一个df，空头也给一个df
    # long_df, short_df = Strategy.before_filter(back_test_df, filter_list)
    long_df, short_df = Strategy.before_filter3(back_test_df, filter_list)
    # 如果使用现货数据，多头选币不受影响，但是空头选币只能在有合约的现货中进行选币
    if if_use_spot:
        short_df = short_df[short_df['tag'] == 'HasSwap']  # 保留有合约的现货

    # ===选币操作
    # 根据选币数量进行多空选币操作
    select_coin = select_long_and_short_coin(long_df, short_df, Strategy.long_select_coin_num,
                                             Strategy.short_select_coin_num,
                                             long_factor=Strategy.long_factor, short_factor=Strategy.short_factor)
    # 如果是现货模式，将现货涨跌幅数据换成合约涨跌幅数据
    if if_use_spot:
        select_coin = transfer_swap(select_coin, df_swap, special_symbol_dict)

    # ===计算资金曲线前的一些操作
    # 针对合约和现货给予不同的手续费率
    select_coin.loc[select_coin['symbol_type'] == 'spot', 'rate'] = spot_c_rate
    select_coin.loc[select_coin['symbol_type'] == 'swap', 'rate'] = swap_c_rate
    select_coin['lvg'] = leverage

    # 保留指定字段
    select_coin = select_coin[['candle_begin_time', 'offset', 'symbol', 'symbol_type', '方向', '每小时涨跌幅', 'fundingRate', 'rate', 'lvg', 'ret_next']]

    # 在选币信息上带上symbol和方向的信息，例如 做空BTC-USDT 会变成 BTC-USDT(-1)
    select_coin['选币'] = select_coin['symbol'] + '(' + select_coin['方向'].astype(str) + ')' + ' '

    # 把选币结果重新排序以及重置索引
    select_coin.sort_values(by='candle_begin_time', inplace=True)
    select_coin.reset_index(drop=True, inplace=True)

    # 输出选币结果（主要是给中性策略查看器用的）
    # select_coin.to_csv(back_test_path + f'{factor_list}_选币结果_{hold_period}_{offset}.csv', index=False, encoding='gbk')
    # print(select_coin)
    #print('选币完成：', time.time() - start_time)

    # ===计算资金曲线
    # =构建回测期间，完整的交易时间列表
    trading_time_list = create_trading_time(benchmark, hold_period, offset, start_date_, end_date_)
    # =将多头和空头拆开计算
    long_df = select_coin[select_coin['方向'] == 1]
    short_df = select_coin[select_coin['方向'] == -1]

    # 本身已多线程
    pool = multiprocessing.Pool(processes=3)
    result1 = pool.apply_async(cal_net_value, (long_df, hold_period, trading_time_list))
    result2 = pool.apply_async(cal_net_value, (short_df, hold_period, trading_time_list))
    result3 = pool.apply_async(calc_swap_pos_net_value,
                               (select_coin, leverage, hold_period, margin_rate, trading_time_list))
    pool.close()
    long_df = result1.get()
    short_df = result2.get()
    swap_net_value = result3.get()

    # long_df = cal_net_value(long_df, hold_period, trading_time_list)
    # short_df = cal_net_value(short_df, hold_period, trading_time_list)
    # swap_net_value = calc_swap_pos_net_value(select_coin, leverage, hold_period, margin_rate, trading_time_list)
    
    #print('计算多空资金曲线 + 合约仓位资金曲线：', time.time() - start_time)

    # 创造空的时间周期表，用于填充不选币的周期
    empty_df = create_empty_data(benchmark, hold_period, offset=offset)
    # 填充缺失数据
    empty_df.update(swap_net_value)
    swap_net_value = empty_df.copy()
    swap_net_value['选币'].fillna(method='ffill', inplace=True)
    swap_net_value.dropna(subset=['选币'], inplace=True)
    swap_net_value['每小时资金曲线'] = swap_net_value['每小时资金曲线'].apply(list)

    # 将多空方向处理后的数据合并起来
    select_coin = pd.concat([long_df, short_df], ignore_index=True)
    #print('多空资金曲线：', time.time() - start_time)
    # ===计算当周期资金曲线
    # 将多空选币的资金曲线合并成当周期内的资金曲线
    group = select_coin.groupby('candle_begin_time')
    # 存储周期内的资金曲线
    merge_df = pd.DataFrame()
    merge_df['选币'] = group['选币'].sum()
    merge_df['offset'] = group['offset'].last()
    # 将多空方向的资金曲线合并起来
    merge_df['每小时资金曲线'] = group['每小时资金曲线'].apply(lambda x: np.array(x).mean(axis=0))
    # 计算周期的涨跌幅
    merge_df['周期涨跌幅'] = merge_df['每小时资金曲线'].apply(lambda x: x[-1] - 1)
    # 通过每小时资金曲线，计算每小时账户资金的涨跌幅
    merge_df['每小时涨跌幅'] = merge_df['每小时资金曲线'].apply(lambda x: list(pd.DataFrame([1] + list(x)).pct_change()[0].iloc[1:]))
    # 计算整个周期多空的平均调仓
    merge_df['调仓比例'] = group['调仓比例'].mean()
    #print('计算当周期资金曲线：', time.time() - start_time)

    # =====计算小时级别的资金曲线
    # 创造空的时间周期表，用于填充不选币的周期
    empty_df = create_empty_data(benchmark, hold_period, offset=offset)
    # 填充缺失数据
    empty_df.update(merge_df)
    merge_df = empty_df.copy()
    merge_df['选币'].fillna(method='ffill', inplace=True)
    merge_df.dropna(subset=['选币'], inplace=True)
    #print('填充当周期资金曲线：', time.time() - start_time)

    # 将合并后的资金曲线，与benchmark合并
    equity = pd.merge(left=benchmark, right=merge_df[['选币', '调仓比例']], on=['candle_begin_time'], how='left', sort=True)
    #需要再次指定范围，不然年化收益会变成全部时间的
    equity = equity[equity['candle_begin_time'] >= pd.to_datetime(start_date_)]
    equity = equity[equity['candle_begin_time'] <= pd.to_datetime(end_date_)]
    # 填充选币数据
    equity['选币'].fillna(method='ffill', inplace=True)
    equity.dropna(subset=['选币'], inplace=True)
    # 将每小时涨跌幅数据，填充到index中
    equity['涨跌幅'] = revise_data_length(merge_df['每小时涨跌幅'].sum(), len(equity))
    equity['涨跌幅'].fillna(value=0, inplace=True)
    # 计算最终每小时净值变化
    equity['资金曲线'] = (equity['涨跌幅'] + 1).cumprod()
    equity['资金曲线'].fillna(value=1, inplace=True)
    # 计算爆仓
    if not swap_net_value.empty:
        equity['合约资金曲线'] = revise_data_length(swap_net_value['每小时资金曲线'].sum(), len(equity), value=1)
        equity.loc[equity['合约资金曲线'] == 0, '是否爆仓'] = 1
        equity.loc[equity['资金曲线'] < margin_rate, '是否爆仓'] = 1
    else:
        equity['是否爆仓'] = 0
    # 这里计算合约仓位爆仓之后，直接将现货仓位也同步归零，实际实盘的时候的会留有部分现货仓位的底仓
    equity['是否爆仓'].fillna(method='ffill', inplace=True)
    equity.loc[equity['是否爆仓'] == 1, '资金曲线'] = 0
    equity.loc[equity['是否爆仓'] == 1, '涨跌幅'] = 0
    equity.loc[equity['是否爆仓'] == 1, '调仓比例'] = 0
    # 保存资金曲线文件
    equity['本周期多空涨跌幅'] = equity['涨跌幅']

    # =====策略评价
    rtn, year_return, month_return = strategy_evaluate(equity)
    
    rtn.loc['offset', 0] = offset
    rtn.loc['持仓周期', 0] = hold_period
    rtn.loc['datename', 0] = datename
    rtn.loc['long', 0] = Strategy.long_select_coin_num
    rtn.loc['short', 0] = Strategy.short_select_coin_num
    index = 1
    for factor_name, if_reverse, parameter_list, weight in factor_list:
        rtn.loc['factor_%s'%index, 0] = factor_name
        rtn.loc['cycle_%s'%index, 0] = parameter_list
        index += 1
    index = 1
    for factor_name, parameter_list, if_reverse, value in filter_list:
        rtn.loc['filter_%s'%index, 0] = factor_name
        rtn.loc['fcycle_%s'%index, 0] = parameter_list
        index += 1
    if not os.path.exists(out_csv_path):
        rtn.T.to_csv(out_csv_path)
    else:
        rtn.T.to_csv(out_csv_path,header =False,mode='a')
    print(factor_list, datename, filter_list, hold_period)
    #return rtn.T

start_time=time.time()
back_test_path = root_path + '/data/回测结果/'
Strategy = __import__('strategy.%s' % strategy_name, fromlist=('',))
offset = 0  # 本次回测设置的offset

output_path=root_path+f'/data/单因子遍历结果/'
if not os.path.exists(output_path):
    os.mkdir(output_path)
factors=['ILLQStd']

#主线程读取文件read_coin2
n_jobs = 90
#遍历的线程，注意线程中有3个线程做cal_net_value
n_jobs_ = 55
parallel_index = 0
datename_list = ['2021','2022','2023']
factor_periods = sorted(list(set(list(range(3, 21, 3)) + list(range(20, 100, 20)) + list(range(100, 1000, 100)) + list(range(1000, 2001, 200)))))
#factor_periods = [3,40,200,600]
filter_periods = factor_periods
filter_name = 'PctChange'    #Bias,PctChange
print('factor_period:', len(factor_periods), 'count:', len(factor_periods)*len(factor_periods)*len(datename_list))
########################注意########################
###1，before_filter函数要根据不同的情况做调整，比如Bias<0，不是按rank排序
for hold_period in ['6H']:   #'6H',
    for factor in factors:
        out_csv_path = output_path + f'{factor}_{filter_name}_{hold_period}_s{Strategy.short_select_coin_num}_l{Strategy.long_select_coin_num}.csv'
        if os.path.exists(out_csv_path):
            os.remove(out_csv_path)
        stime = time.time()
        df2 = read_coin2(root_path, hold_period, [factor,filter_name], if_use_spot, n_jobs, offset)
        print(time.strftime("%Y-%m-%d %H:%M:%S",time.localtime()), 'read coin done!', factor, time.time()-stime)
        para_lists = []
        for datename in datename_list:
            start_date_ = datename + '-01-01'
            end_date_ = datename + '-12-31'
            df = df2.copy()
            df = df[df['candle_begin_time'] >= pd.to_datetime(start_date_)]
            df = df[df['candle_begin_time'] <= pd.to_datetime(end_date_)]
            for factor_cycle in factor_periods:
                #根据因子，自行指定True或False，默认为True
                factor_list = [(factor, True, factor_cycle, 1),]
                for filter_cycle in filter_periods:
                    filter_list = [(filter_name, filter_cycle, True, 0.8),]
                    para_lists.append([datename, factor_list, filter_list])
            #如果内存不够，可以再加一个缩进，即每个facto_cycle进行一轮
            print(factor, datename, 'parallel_count', len(para_lists))
            Parallel(n_jobs=n_jobs_)(delayed(handle_data)(out_csv_path, df, datename_2, factor_list_2, filter_list_2) for datename_2, factor_list_2, filter_list_2 in para_lists)
            para_lists=[]
            print(time.strftime("%Y-%m-%d %H:%M:%S",time.localtime()), 'datename done!', factor, datename, time.time()-stime)
        print('factor done!', filter_name, factor, hold_period, time.time()-stime)
    print('hold_period done!', filter_name, factor, hold_period, time.time()-stime)
print('all done!')